from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
import isaacgym
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_env import *

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback


# # get console input params
# import argparse
# parser = argparse.ArgumentParser()
# parser.add_argument('--gen', required=True, help='generation')
# parser.add_argument('--dec', required=True, help='decendent')
# opt = parser.parse_args()
# print(opt)


LOGDIR = "./training_logs"

# m_env = CustomEnv(seed_pth = "./gen_0_dec_2" + ".json", headless_ = False)      # 146.10514254070014
# m_env = CustomEnv(seed_pth = "./gen_1_dec_4" + ".json", headless_ = False)      # 167.55190241091634
# m_env = CustomEnv(seed_pth = "./gen_2_dec_6" + ".json", headless_ = False)      # 195.69338548700745
# m_env = CustomEnv(seed_pth = "./gen_3_dec_2" + ".json", headless_ = False)      # 203.01231236390612
# m_env = CustomEnv(seed_pth = "./gen_4_dec_1" + ".json", headless_ = False)      # 199.56112463758984
# m_env = CustomEnv(seed_pth = "./gen_5_dec_7" + ".json", headless_ = False)      # 206.61725791020106
# m_env = CustomEnv(seed_pth = "./gen_6_dec_7" + ".json", headless_ = False)      # 216.4248953390794
# m_env = CustomEnv(seed_pth = "./gen_7_dec_3" + ".json", headless_ = False)      # 214.9650354474695
# m_env = CustomEnv(seed_pth = "./gen_8_dec_5" + ".json", headless_ = False)      # 209.6587573199749
# m_env = CustomEnv(seed_pth = "./gen_9_dec_2" + ".json", headless_ = False)      # 213.15502561786212
# m_env = CustomEnv(seed_pth = "./gen_10_dec_1" + ".json", headless_ = False)     # 202.23301456205942
# m_env = CustomEnv(seed_pth = "./gen_11_dec_5" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_12_dec_1" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_13_dec_4" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_14_dec_6" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_15_dec_0" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_16_dec_5" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_17_dec_1" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_18_dec_1" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_19_dec_5" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_20_dec_3" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_21_dec_6" + ".json", headless_ = False)     # 205.03441441394037
m_env = CustomEnv(seed_pth = "./gen_22_dec_4" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv(seed_pth = "./gen_23_dec_5" + ".json", headless_ = False)     # 205.03441441394037
# m_env = CustomEnv( headless_ = False)
observation = m_env.reset()

m_env.headless = False

print("num_envs:",m_env.num_envs)

pcg_worker = PCGWorker(9,9)

print("fitness --------------- ",pcg_worker.fitness(m_env.seeds[0]))

num_tile = 0
for i in range(0,9):
    for j in range(0,9):
        tile_ = m_env.seeds[0].wave_oriented[j*9+i][0][0]
        rot = m_env.seeds[0].wave_oriented[j*9+i][0][1]
        if tile_ == 2:
            num_tile += 1
            
print("num_tile --------------- ",num_tile)

action = -1
while True:
    
    observation, reward, done, info = m_env.step(action)

    # m_env.PCGWorker_.render(m_env.seeds[0])

    # print(observation)
    # print(reward)
    # print(done)
    # print(info)

    if done:
        observation = m_env.reset()

    interaction = m_env.render()

    if interaction == -2:
        m_env.close()

    if interaction == 119:  # w
        action = 0
    elif interaction == 115:    # s
        action = 1
    elif interaction == 97:     # a
        action = 2
    elif interaction == 100:    # d
        action = 3
    elif interaction == 107:    # ccw
        action = 4
    elif interaction == 108:    # cw
        action = 5
    # elif interaction == 111:    # break
    #     action = 6
    elif interaction == 114:
        m_env.reset()
    else:
        action = -1

